from websocket import create_connection
import json
from .language_model import LanguageModel

class PentestMuseLanguageModel(LanguageModel):
    def __init__(self, credentials):
        self.token = credentials
        self.url = "wss://api.pentestmuse.ai/cli/chat_socket"

    def generate(self, messages):
        ws = create_connection(self.url)

        # Send the messages to the server
        ws.send(json.dumps(
            {
                "messages": messages,
                "token": self.token
            }
        ))

        while True:
            result = ws.recv()
            if not result:
                break

            result_json = json.loads(result)
            if 'error' in result_json:
                print(f"Error: {result_json['error']}")
                break

            if 'content' in result_json:
                yield result_json['content']

            if result_json.get('isEndOfStream'):
                break